
%% parameters to set ahead of time:
doGcampBloodCorrection = true;      %sets whether you do a bloodflow correction to the gcamp values
skipIOStracking = false;            %set to false to skip IOS tracking; can be useful if just want to re-calculate the image processing stuff or  gcamp values
savedata = true;                   %set to true to automatically save the results
downsamplegcamp = 4;             %allow pixel binning for gcamp to make things go faster/take less memory
%% load in previously selected referencing information:
%[fn,fd] = uigetfile('*.mat','select referencing information .mat file');
cd(fd);
load(fn);

%% load in vessel mask:
load(vesselmaskname,'vesselMasks','avios','refim_windowmask','micronsPerPixel')     %avios is what we're going to be registering everything to; there is ONE of these per mouse 
refim = avios;              %need to make a reference ios image for each session

%% load in background:
cd(bgdirectory);
try
    bg = loadRawStack_v1(rootname,localimagesize);
catch
    bg = loadRawStack_v1(rootname,localimagesize,[],[],[],'uint8');
    disp('FAILED TO LOAD 16 BIT - RETRYING WITH 8')
end
bg_gcamp = bg(:,:,1:2:end);
bg_ios = bg(:,:,2:2:end);
load(bgmetaname,'movindicator','rectPositions','framesPerStim','stimparams')
bg_gcamp = reshape(bg_gcamp,size(bg_gcamp,1),size(bg_gcamp,2),framesPerStim/2,length(movindicator));
bg_ios = reshape(bg_ios,size(bg_ios,1),size(bg_ios,2),framesPerStim/2,length(movindicator));

%an inelegant way of dealing with the moving stim experiments - currently
%not generalizable
if exist('stimtype','var') && strcmp(stimtype,'moving')
    load(bgmetaname,'xp','yp');
    stimStartPos = zeros(length(movindicator),2);
    stimStartPos(movindicator == 1,:) = [xp(1),yp(1)];
    stimStartPos(movindicator == 2,:) = [xp(end),yp(end)];
    stimStartPos(movindicator == 3,:) = [.001,.001];
    bguid = [stimStartPos';movindicator];
    stimsizes = zeros(size(movindicator));
    stimsizes(movindicator~=3) = stimparams{1}.circlesize;
    bguid = [bguid;stimsizes];
else
    bguid = [reshape(rectPositions,2,size(rectPositions,3));movindicator];

    stimsizes = zeros(size(stimparams));
    for n = 1:length(stimparams)
        switch stimparams{n}.stimtype
            case 'circle'
                stimsizes(n) = stimparams{n}.circlesize;
            case 'blank'
                stimsizes(n) = 0;
            case 'bpFilteredNoise'
                stimsizes(n) = stimparams{n}.circlesize;
        end
    end
    
    temp = stimsizes(movindicator);
    bguid = [bguid;temp];
end
[uid_bg,~,bgID_idx] = unique(bguid','rows');

clear stimparams rectPositions movindicator xp yp
%% load in metadata
load(metaname,'movindicator','rectPositions','framesPerStim','preStimTime','totalFrameTime','stimparams');
if exist('stimtype','var') && strcmp(stimtype,'moving')
    load(metaname,'xp','yp');
    stimStartPos = zeros(length(movindicator),2);
    stimStartPos(movindicator == 1,1) = xp(1);
    stimStartPos(movindicator == 1,2) = yp(1);
    stimStartPos(movindicator == 2,1) = xp(end);
    stimStartPos(movindicator == 2,2) = yp(end);
    stimStartPos(movindicator == 3,:) = .001;
    stimuid = [stimStartPos';movindicator];
    stimsizes = zeros(size(movindicator));
    stimsizes(movindicator~=3) = stimparams{1}.circlesize;
    stimuid = [stimuid;stimsizes];
else
    stimuid = [reshape(rectPositions,2,size(rectPositions,3));movindicator];
    stimsizes = zeros(size(stimparams));
    for n = 1:length(stimparams)
        switch stimparams{n}.stimtype
            case 'circle'
                stimsizes(n) = stimparams{n}.circlesize;
            case 'blank'
                stimsizes(n) = 0;
            case 'bpFilteredNoise'
                stimsizes(n) = stimparams{n}.circlesize;
        end
    end
    temp = stimsizes(movindicator);
    stimuid = [stimuid;temp];
end
[uid,~,stimID_idx] = unique(stimuid','rows');
totalFrameTime = totalFrameTime/1000;       %convert into seconds

time = (1:(framesPerStim/2))*totalFrameTime;
prestimframes = 2:sum(time<preStimTime);
stimframes = (prestimframes(end)+1):sum(time<(preStimTime + stimparams{1}.totalStimLength));
iosResponseFrames = find(time>(preStimTime+1),1):find(time>(preStimTime + stimparams{1}.totalStimLength+1),1);
%% will want to get ONE very good tranformation matrix to map the data from the current session to the master session
% to do this you need to load in some ios data from the current session,
% make an average image; do NOT need to do this if this is the reference
% session:
if ~isReferenceSession
    cd(movdirectory);
    curidx = 1:(max(prestimframes)*2);
    localavios = loadRawStack_v1(rootname,localimagesize,curidx);
    localavios = localavios(:,:,2:2:end);
    localavios = mean(localavios,3);
    A = localavios;A(~windowmask) = median(A(:));
    B = refim; B(~refim_windowmask) = median(B(:));
    if exist('sessionRegister','var')
        %use the pre-made session registration
    else
        sessionRegister = registerMultipleWidefieldImages_v1_2(A,B);
    end
    localavios_transformed = imwarp(localavios,sessionRegister.Transformation,'OutputView',imref2d([size(refim,1),size(refim,2)]));
    figure; imshowpair(refim,localavios_transformed,"checkerboard");
    clear A B
else
    localavios = avios;
end
%% loop through movies and do the tracking (for .raw formatted image sequences):
cd(movdirectory);
totalframes = framesPerStim*length(movindicator);
traj = [];
gcampmov_full = [];
parenchymaMask = false(size(bg_ios,1),size(bg_ios,2),length(movindicator));
roughvesselmask = false(size(bg_ios,1),size(bg_ios,2),length(movindicator));

avgParenchymaIOS_normalized = zeros(size(parenchymaMask));
for n = 1:length(movindicator)
    % select only the frames corresponding to the given movie:
    curidx = ((n-1)*framesPerStim + 1):(n*framesPerStim);
    mov = loadRawStack_v1(rootname,localimagesize,curidx);
    iosmov = mov(:,:,2:2:end);
    curgcampmov = mov(:,:,1:2:end);
    
    %subtract appropriate background:
    curbg = bgID_idx==stimID_idx(n);
    curgcampmov = curgcampmov-bg_gcamp(:,:,:,curbg);
    iosmov = iosmov-bg_ios(:,:,:,curbg);
    
    if ~isReferenceSession
        [iosmov,curgcampmov,curtraj] = trackOneWidefieldMovie(iosmov,curgcampmov,windowmask,localavios,refim,vesselMasks, sessionRegister,doGcampBloodCorrection,prestimframes,skipIOStracking);
    else
        [iosmov,curgcampmov,curtraj] = trackOneWidefieldMovie(iosmov,curgcampmov,windowmask,localavios,[],vesselMasks,[],doGcampBloodCorrection,prestimframes,skipIOStracking);
    end
    [parenchymaMask(:,:,n),roughvesselmask(:,:,n)] = getParenchymaPixels(mean(iosmov,3),windowmask);
    
    for n2 = 1:size(curgcampmov,3)
        temp = curgcampmov(:,:,n2);
        temp(roughvesselmask(:,:,n)) = nan;
        temp(~windowmask) = 0 ;
        curgcampmov(:,:,n2) = temp;
        
        temp = iosmov(:,:,n2);
        temp(~parenchymaMask(:,:,n)) = nan;
        temp(~windowmask) = 0;
        iosmov(:,:,n2) = temp;
    end
    temp = mean(iosmov(:,:,prestimframes),3,'omitnan');
    temp = iosmov./repmat(temp,1,1,size(iosmov,3));
    avgParenchymaIOS_normalized(:,:,n) = mean(temp(:,:,iosResponseFrames),3,'omitnan');
    
    %downsample gcamp
    if downsamplegcamp > 1
        for n2 = 1:size(curgcampmov,3)
            temp = curgcampmov(:,:,n2);
            temp = binimage(temp,downsamplegcamp);
            if n2 == 1
                if isempty(gcampmov_full)
                    idx = 1;
                else
                    idx = size(gcampmov_full,3) + 1;
                end
                gcampmov_full = cat(3,gcampmov_full,zeros(size(temp,1),size(temp,2),size(curgcampmov,3)));
            else
                idx = idx + 1;
            end
            gcampmov_full(:,:,idx) = temp;
        end
    else
        gcampmov_full = cat(3,gcampmov_full,curgcampmov);
    end
    traj = [traj,curtraj];
    disp(['tracked movie # ',num2str(n)])
end
numTrackedSegs = size(traj(1).vesselCent,1);
disp('FINISHED TRACKING DILATIONS')
%% calculate baseline levels for vessel total intensity:
clear iosmov mov curgcampmov curtraj
baselineVesselInten = zeros(numTrackedSegs,length(traj));
if ~skipIOStracking
    for n = 1:length(traj)
        baselineVesselInten(:,n) = mean(traj(n).vesselTotalInten(:,prestimframes),2,'omitnan');
    end
    baselineVesselInten = median(baselineVesselInten,2,'omitnan');
end
%% calculate  baseline levels for surround total intensity:
baselineSurroundInten = zeros(numTrackedSegs,length(traj));
if ~skipIOStracking
    for n = 1:length(traj)
        baselineSurroundInten(:,n) = mean(traj(n).surroundInten(:,prestimframes),2,'omitnan');
    end
    baselineSurroundInten = median(baselineSurroundInten,2,'omitnan');
end
    
%% run processing on gcamp movie AFTER the fact to allow for global baseline correction:
clear iosmov curgcampmov curtraj
downsampledwindowmask = binimage(windowmask,downsamplegcamp);
downsampledwindowmask = downsampledwindowmask > 0;
[gcampmov_norm,baselineGcamp,baselineStdGcamp,GcampNormalizationRange] = gcampProcess_v1_0(gcampmov_full,micronsPerPixel*downsamplegcamp,GcampSmoothSize,framesPerStim,prestimframes,totalFrameTime,downsampledwindowmask);


gcampfloor = .2;
gcampmov_norm(gcampmov_norm<gcampfloor) = 0;
gcampmov_norm(isnan(gcampmov_norm)) = 0;
disp('gcamp movie image processing done')
%% gcamp distance calculation:
gcampDistanceScaling = 400;         %in microns
localgcamplevel = zeros(numTrackedSegs,size(gcampmov_norm,3),size(gcampmov_norm,4));
localgcampdistance = zeros(numTrackedSegs,size(gcampmov_norm,3),size(gcampmov_norm,4));
for n = 1:size(gcampmov_norm,4)
    for n2 = 1:size(gcampmov_norm,3)
        temp = normalizedNeuralActivityDistance_v1_2(gcampmov_norm(:,:,n2,n),traj(1).vesselCent/downsamplegcamp,micronsPerPixel*downsamplegcamp,gcampDistanceScaling,'measured',false);
        temp2 = log(temp.distanceWeightedActivity);
        temp2(temp2<0) = 0;
        localgcamplevel(:,n2,n) = temp2;
        localgcampdistance(:,n2,n) = temp.distance2activity;
    end
    disp(n)
end
disp('gcamp distance calculations done')

%% generate average gcamp images for storing:
avgNormalizedGcampImages = zeros(size(gcampmov_norm,1),size(gcampmov_norm,2),size(gcampmov_norm,4));
for n = 1:size(gcampmov_norm,4)
    avgNormalizedGcampImages(:,:,n) = mean(gcampmov_norm(:,:,stimframes,n),3);
end

%% save data so it's not lost
cd(savedirectory)
if savedata
    if skipIOStracking
        temp = strrep(savename,'.mat','_noIOStracking.mat');
        save(temp,'localgcamplevel','localgcampdistance','gcampDistanceScaling','avgNormalizedGcampImages','stimID_idx','uid','parenchymaMask','roughvesselmask','avgParenchymaIOS_normalized',...
            'baselineGcamp','baselineStdGcamp','GcampNormalizationRange','micronsPerPixel','downsamplegcamp','avios');
    else
        save(savename,'localgcamplevel','localgcampdistance','gcampDistanceScaling','baselineVesselInten','traj','avgNormalizedGcampImages','stimID_idx','uid','parenchymaMask','roughvesselmask','baselineSurroundInten','avgParenchymaIOS_normalized',...
            'baselineGcamp','baselineStdGcamp','GcampNormalizationRange','micronsPerPixel','downsamplegcamp','avios');
    end
end



